import { useQueryClient } from '@tanstack/react-query';
import React, {
    createContext,
    useState,
    useEffect,
    useContext,
    useCallback,
} from 'react';

interface metric {
    name: string;
    value: any;

}

interface WebSocketData {
    id: string;
    online: boolean;
    metrics: metric[];
    timestamp: number;
}

interface WebSocketContextType {
    socket: WebSocket | null;
}

const WebSocketContext = createContext<WebSocketContextType | undefined>(
    undefined
);

export const WebSocketProvider: React.FC<{ children: React.ReactNode }> = ({
    children,
}) => {
    const [socket, setSocket] = useState<WebSocket | null>(null);
    const queryClient = useQueryClient();

    const handleWebSocketMessage = useCallback((event: MessageEvent) => {
        try {
            const message: WebSocketData = JSON.parse(event.data);
            queryClient.setQueryData(['state', message.id], message);
        } catch (error) {
            console.error('Error parsing WebSocket message:', error);
        }
    }, [queryClient]);

    useEffect(() => {
        const ws = new WebSocket('ws://localhost:3333/ws');

        ws.onopen = () => {
            console.log('WebSocket connected');
            setSocket(ws);
        };

        ws.onmessage = handleWebSocketMessage;

        ws.onclose = () => {
            console.log('WebSocket disconnected');
            setSocket(null);
        };

        return () => {
            if (ws) {
                ws.close();
            }
        };
    }, [handleWebSocketMessage]);

    return (
        <WebSocketContext.Provider value={{ socket }}>
            {children}
        </WebSocketContext.Provider>
    );
};

